package com.ahdms.es.util;

import com.ahdms.es.gm.algorithm.SM4Engine;
import org.bouncycastle.crypto.params.KeyParameter;

import java.io.ByteArrayOutputStream;

public class SM4Utils {

    public static int littleEndianToShort(byte[] bs, int off) {
        int n = bs[off] & 0xff;
        n |= (bs[++off] & 0xff) << 8;
        return n;
    }

    public static void shortToLittleEndian(int n, byte[] bs, int off) {
        bs[off] = (byte) (n);
        bs[++off] = (byte) (n >>> 8);
    }

    /**
     * 按每63K字节对原文进行分组
     *
     * @param plainText
     * @return
     */
    private static byte[][] groupingPlainText(byte[] plainText) {

        int maxDataLengthPerGroup = 64 * 1024 - 1;
        int groupSize = (int) Math.ceil((double) plainText.length / maxDataLengthPerGroup);
        byte[][] groupedData = new byte[groupSize][];

        int left = plainText.length;
        int j = 0;
        int curDataLen = 0;
        do {
            if (left > maxDataLengthPerGroup) {
                curDataLen = maxDataLengthPerGroup;
            } else {
                curDataLen = left;
            }

            int paddingLength = 0;
            if (0 != (3 + curDataLen) % 16) {
                paddingLength = 16 - ((3 + curDataLen) % 16);
            }

            groupedData[j] = new byte[3 + curDataLen + paddingLength];

            left -= curDataLen;
            shortToLittleEndian(curDataLen, groupedData[j], 0);
            System.arraycopy(plainText, j * maxDataLengthPerGroup, groupedData[j], 2, curDataLen);

            groupedData[j][2 + curDataLen] = (byte) 0x80;

            j++;

        } while (left > 0);

        return groupedData;
    }

    private static void groupEncryptCBC(SM4Engine sm4Engine, byte[] iv, byte[] group, byte[] out) throws Exception {
        if (iv == null || iv.length != 16) {
            throw new IllegalArgumentException("iv error!");
        }

        int blockSize = sm4Engine.getBlockSize();
        int blockNum = group.length / blockSize;
        byte[] tmpIn = new byte[blockSize];
        byte[] tmpOut = new byte[blockSize];
        byte[] tmpIv = new byte[blockSize];
        int off = 0;
        System.arraycopy(iv, 0, tmpIv, 0, blockSize);

        for (int i = 0; i < blockNum; i++) {
            off = i * blockSize;
            System.arraycopy(group, off, tmpIn, 0, blockSize);
            blockProcessCBC(sm4Engine, true, tmpIv, tmpIn, tmpOut);
            System.arraycopy(tmpOut, 0, tmpIv, 0, blockSize);
            System.arraycopy(tmpOut, 0, out, off, blockSize);
        }
    }

    private static void blockProcessCBC(SM4Engine sm4Engine, boolean forEncryption, byte[] iv, byte[] in, byte[] out) throws Exception {
        if (iv == null || iv.length != 16) {
            throw new IllegalArgumentException("iv error!");
        }

        byte[] tmp = new byte[16];
        if (forEncryption) {
            for (int i = 0; i < 16; i++) {
                tmp[i] = ((byte) (in[i] ^ iv[i]));
            }
            sm4Engine.processBlock(tmp, 0, out, 0);
        } else {
            sm4Engine.processBlock(in, 0, tmp, 0);
            for (int i = 0; i < 16; i++) {
                out[i] = ((byte) (tmp[i] ^ iv[i]));
            }
        }
    }

    public static byte[] encryptDataCBC(byte[] plainText, byte[] key, byte[] iv) throws Exception {
        boolean forEncryption = true;

        SM4Engine sm4Engine = new SM4Engine();
        KeyParameter keyParam = new KeyParameter(key);
        sm4Engine.init(forEncryption, keyParam);

        int blockSize = sm4Engine.getBlockSize();
        byte[][] groupedData = groupingPlainText(plainText);
        int groupSize = groupedData.length;
        ByteArrayOutputStream bous = new ByteArrayOutputStream();

        byte[] tmpIv = new byte[blockSize];
        System.arraycopy(iv, 0, tmpIv, 0, blockSize);

        for (int i = 0; i < groupSize; i++) {
            byte[] tmpOut = new byte[groupedData[i].length];
            groupEncryptCBC(sm4Engine, tmpIv, groupedData[i], tmpOut);
            System.arraycopy(tmpOut, tmpOut.length - blockSize, tmpIv, 0, blockSize);
            bous.write(tmpOut, 0, tmpOut.length);
        }

        byte[] out = bous.toByteArray();
        bous.close();
        return out;
    }

    private static void blockProcessECB(SM4Engine sm4Engine, byte[] in, byte[] out) throws Exception {
        sm4Engine.processBlock(in, 0, out, 0);
    }

    private static void groupEncryptECB(SM4Engine sm4Engine, byte[] group, byte[] out) throws Exception {

        int blockSize = sm4Engine.getBlockSize();
        int blockNum = group.length / blockSize;
        byte[] tmpIn = new byte[blockSize];
        byte[] tmpOut = new byte[blockSize];
        int off = 0;

        for (int i = 0; i < blockNum; i++) {
            off = i * blockSize;
            System.arraycopy(group, off, tmpIn, 0, blockSize);
            blockProcessECB(sm4Engine, tmpIn, tmpOut);
            System.arraycopy(tmpOut, 0, out, off, blockSize);
        }
    }

    public static byte[] encryptDataECB(byte[] plainText, byte[] key) throws Exception {
        boolean forEncryption = true;

        SM4Engine sm4Engine = new SM4Engine();
        KeyParameter keyParam = new KeyParameter(key);
        sm4Engine.init(forEncryption, keyParam);

        byte[][] groupedData = groupingPlainText(plainText);
        int groupSize = groupedData.length;
        ByteArrayOutputStream bous = new ByteArrayOutputStream();

        for (int i = 0; i < groupSize; i++) {
            byte[] tmpOut = new byte[groupedData[i].length];
            groupEncryptECB(sm4Engine, groupedData[i], tmpOut);
            bous.write(tmpOut, 0, tmpOut.length);
        }

        byte[] out = bous.toByteArray();
        bous.close();
        return out;
    }

    public static byte[] decryptDataCBC(byte[] cipherText, byte[] key, byte[] iv) throws Exception {
        boolean forEncryption = false;

        SM4Engine sm4Engine = new SM4Engine();
        KeyParameter keyParam = new KeyParameter(key);
        sm4Engine.init(forEncryption, keyParam);

        int blockSize = sm4Engine.getBlockSize();
        int blockNum = cipherText.length / blockSize;

        byte[] tmpBlock = new byte[blockSize];
        byte[] tmpOut = new byte[blockSize];
        byte[] tmpIv = new byte[blockSize];
        System.arraycopy(iv, 0, tmpIv, 0, blockSize);

        ByteArrayOutputStream bous = new ByteArrayOutputStream();

        for (int i = 0; i < blockNum; i++) {
            System.arraycopy(cipherText, i * blockSize, tmpBlock, 0, blockSize);
            blockProcessCBC(sm4Engine, forEncryption, tmpIv, tmpBlock, tmpOut);
            System.arraycopy(tmpBlock, 0, tmpIv, 0, blockSize);
            bous.write(tmpOut, 0, blockSize);
        }

        byte[] out = parsePlainText(bous.toByteArray());
        bous.close();

        return out;
    }

    private static byte[] parsePlainText(byte[] in) throws Exception {
        ByteArrayOutputStream bous = new ByteArrayOutputStream();

        int tmpDataLen = 0;
        for (int i = 0; i < in.length; ) {
            if (0 == tmpDataLen) {
                tmpDataLen = littleEndianToShort(in, i);
                i += 2;
            }

            bous.write(in, i, tmpDataLen);
            i += tmpDataLen;  //data

            tmpDataLen = 0;
            i++; //0x80

            if (0 != i % 16) {
                i += (16 - i % 16);
            }
        }

        byte[] out = bous.toByteArray();
        bous.close();
        return out;
    }

    public static byte[] decryptDataECB(byte[] cipherText, byte[] key) throws Exception {
        boolean forEncryption = false;

        SM4Engine sm4Engine = new SM4Engine();
        KeyParameter keyParam = new KeyParameter(key);
        sm4Engine.init(forEncryption, keyParam);

        int blockSize = sm4Engine.getBlockSize();
        int blockNum = cipherText.length / blockSize;

        byte[] tmpBlock = new byte[blockSize];
        byte[] tmpOut = new byte[blockSize];

        ByteArrayOutputStream bous = new ByteArrayOutputStream();

        for (int i = 0; i < blockNum; i++) {
            System.arraycopy(cipherText, i * blockSize, tmpBlock, 0, blockSize);
            blockProcessECB(sm4Engine, tmpBlock, tmpOut);
            bous.write(tmpOut, 0, blockSize);
        }

        byte[] out = parsePlainText(bous.toByteArray());
        bous.close();

        return out;
    }

    private static byte[] padPlainText(byte[] plainText, int blockSize) {
        int origTextLen = plainText.length;
        int paddingLen = 0;
        if (0 == origTextLen % blockSize) {
            paddingLen = blockSize;
        } else {
            paddingLen = blockSize - origTextLen % blockSize;
        }

        byte[] paddedData = new byte[origTextLen + paddingLen];
        System.arraycopy(plainText, 0, paddedData, 0, origTextLen);
        for (int i = 0; i < paddingLen; i++) {
            paddedData[origTextLen + i] = (byte) paddingLen;
        }
        return paddedData;
    }

    private static byte[] unPadPlainText(byte[] origPlainText, int blockSize) {
        int origPlainTextLen = origPlainText.length;

        boolean paddingFlag = true;
        if (0 == origPlainTextLen % blockSize) {
            for (int i = 0; i < blockSize; i++) { //判断最后一个block是否是填充的值
                if (origPlainText[origPlainTextLen - i - 1] != (byte) blockSize) {
                    paddingFlag = false;
                    break;
                }
            }
        }

        if (paddingFlag) {
            int paddingLen = origPlainText[origPlainTextLen - 1];
            int origTextLen = origPlainTextLen - paddingLen;

            byte[] unpadData = new byte[origTextLen];
            System.arraycopy(origPlainText, 0, unpadData, 0, origTextLen);

            return unpadData;
        } else {
            return origPlainText;
        }
    }

    public static byte[] encryptDataECBNoGroup(byte[] plainText, byte[] key) throws Exception {
        boolean forEncryption = true;

        SM4Engine sm4Engine = new SM4Engine();
        KeyParameter keyParam = new KeyParameter(key);
        sm4Engine.init(forEncryption, keyParam);

        int blockSize = sm4Engine.getBlockSize();

        byte[] paddedData = padPlainText(plainText, blockSize);

        int blockNum = paddedData.length / blockSize;
        byte[] tmpIn = new byte[blockSize];
        byte[] tmpOut = new byte[blockSize];
        int off = 0;

        ByteArrayOutputStream bous = new ByteArrayOutputStream();
        for (int i = 0; i < blockNum; i++) {
            off = i * blockSize;
            System.arraycopy(paddedData, off, tmpIn, 0, blockSize);
            blockProcessECB(sm4Engine, tmpIn, tmpOut);
            bous.write(tmpOut, 0, blockSize);
        }

        byte[] out = bous.toByteArray();
        bous.close();

        return out;
    }

    public static byte[] decryptDataECBNoGroup(byte[] cipherText, byte[] key) throws Exception {
        boolean forEncryption = false;

        SM4Engine sm4Engine = new SM4Engine();
        KeyParameter keyParam = new KeyParameter(key);
        sm4Engine.init(forEncryption, keyParam);

        int blockSize = sm4Engine.getBlockSize();
        int blockNum = cipherText.length / blockSize;

        byte[] tmpBlock = new byte[blockSize];
        byte[] tmpOut = new byte[blockSize];

        ByteArrayOutputStream bous = new ByteArrayOutputStream();

        for (int i = 0; i < blockNum; i++) {
            System.arraycopy(cipherText, i * blockSize, tmpBlock, 0, blockSize);
            blockProcessECB(sm4Engine, tmpBlock, tmpOut);
            bous.write(tmpOut, 0, blockSize);
        }

        byte[] out = unPadPlainText(bous.toByteArray(), blockSize);
        bous.close();

        return out;
    }
}
